package com.liuqi.openai.chain;

import com.liuqi.openai.chain.result.ChainResult;
import com.liuqi.openai.core.chat.Tool;
import com.liuqi.openai.data.message.UserMessage;
import com.liuqi.openai.memory.ChatMemory;
import com.liuqi.openai.model.OpenAiChatModel;
import com.liuqi.openai.tool.CustomTool;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static com.liuqi.openai.memory.chat.MessageWindowChatMemory.withMaxMessages;
import static com.liuqi.openai.util.ValidationUtil.ensureNotBlank;
import static com.liuqi.openai.util.ValidationUtil.ensureNotNull;

/**
 * 用于与指定的{@link OpenAiChatModel}进行交互，同时保持对话记忆的链。
 *
 * 包括一个默认{@link ChatMemory}(最多包含10条消息的消息窗口)，可以被覆盖
 *
 * 实现了工具调用逻辑流程，并将工具调用结果发送给Ai，由AI生成结果 {@link CustomTool}
 *
 * @author liuqi
 * @date 2025/7/20
 **/
public abstract class AbstractOpenAiChain<Input> implements OpenAiChain<Input> {
    private final static Logger logger = LoggerFactory.getLogger(AbstractOpenAiChain.class);

    /**
     * 对话链ID
     */
    protected final String id;

    /**
     * 对话模型
     */
    protected final OpenAiChatModel chatModel;

    /**
     * 对话记忆
     */
    private final ChatMemory chatMemory;

    /**
     * 工具调用
     */
    private final List<? extends CustomTool> customTools;

    /**
     * key: tool name
     * value: CustomTool
     */
    private final Map<String, CustomTool> customToolMap = new HashMap<>();

    /**
     * 当前对话的变量数据, 对话结束后, 需要清空
     */
    private final Map<String, String> currentChatVariables = new HashMap<>();

    /**
     * 当前对话的执行器, 每当对话结束需要赋予空值
     */
    private volatile ChainExecutor currentChainExecutor = null;

    public AbstractOpenAiChain(String id,
                               OpenAiChatModel chatModel,
                               ChatMemory chatMemory,
                               List<? extends CustomTool> customTools) {
        this.id = ensureNotBlank(id, "id");
        this.chatModel = ensureNotNull(chatModel, "chatModel");
        this.chatMemory = chatMemory == null ? withMaxMessages(10) : chatMemory;
        this.customTools = customTools;

        logger.info("create OpenAiChain " + this.getClass().getName());

        if (chatModel.getTools() != null) {
            throw new IllegalArgumentException("setting tools in model is not supported when using the open ai chain module.");
        }

        if (this.customTools != null) {
            logger.info("chain settings tools, tools.size = {}", this.customTools.size());

            List<Tool> modelTools = new ArrayList<>(this.customTools.size());
            this.customTools.stream()
                    .forEach(customTool -> {
                        modelTools.add(customTool.chatTool());
                        customToolMap.put(customTool.toolName(), customTool);
                    });
            chatModel.setTools(modelTools);
        }
    }

    @Override
    public OpenAiChain<Input> injectCurrentChatVariables(Map<String, String> variables) {
        if (variables != null) {
            currentChatVariables.putAll(variables);
        }
        return this;
    }

    @Override
    public synchronized ChainResult execute(Input message) throws OpenAiChainException {
        return execute(message, null);
    }

    @Override
    public synchronized ChainResult execute(Input input, ChainStreamingHandler handler) throws OpenAiChainException {
        if (currentChainExecutor != null) {
            throw new OpenAiChainException("there are currently unfinished chat");
        }

        if (logger.isTraceEnabled()) {
            logger.trace("OpenAiChain chat content: " + input);
        }

        try {
            // 创建对话执行器
            currentChainExecutor = new ChainExecutor(
                    chatModel,
                    chatMemory,
                    customToolMap,
                    currentChatVariables,
                    handler
            );

            return doExecute(input, currentChainExecutor);
        } catch (Throwable t) {
            logger.error("OpenAiChain Error. ", t);
            if (t instanceof OpenAiChainException) {
                throw (OpenAiChainException) t;
            } else {
                throw new OpenAiChainException("OpenAiChain execute error.", t);
            }
        } finally {
            // 当前对话置空
            currentChainExecutor = null;

            // 对话变量清空
            currentChatVariables.clear();
        }
    }

    @Override
    public void stop() {
        if (currentChainExecutor != null) {
            currentChainExecutor.stop();
        }
    }

    /**
     * 由子类执行, 调用 {@link ChainExecutor#execute(UserMessage)} 发起openai请求
     *
     * @param input
     * @return
     * @throws Throwable
     */
    protected abstract ChainResult doExecute(Input input, ChainExecutor executor) throws Throwable;

}
